import asyncio
import json
from datetime import datetime
from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Query
from sqlalchemy.orm import Session
from sqlalchemy import func
from jose import jwt
from api.database import SessionLocal
from api.model.user import User
from api.model.email_account import EmailAccount
from api.model.email import Email
from api.model.email_folder import EmailFolder
from api.controller.auth import get_current_user_from_token, SECRET_KEY, ALGORITHM, ACCESS_TOKEN_EXPIRE_MINUTES, create_access_token
from api.utils.logger import get_logger
from datetime import timedelta

# 初始化日志
logger = get_logger("ws")

router = APIRouter()

class ConnectionManager:
    def __init__(self):
        self.active_connections: dict[int, WebSocket] = {}
        self.last_unread_counts: dict[int, dict] = {}  # 缓存每个用户上次推送的未读数

    async def connect(self, websocket: WebSocket, user_id: int):
        await websocket.accept()
        self.active_connections[user_id] = websocket
        self.last_unread_counts[user_id] = None  # 初始化为None,标记首次连接

    def disconnect(self, user_id: int):
        if user_id in self.active_connections:
            del self.active_connections[user_id]
        if user_id in self.last_unread_counts:
            del self.last_unread_counts[user_id]  # 清理缓存
    
    def _counts_changed(self, old_counts: dict, new_counts: dict) -> bool:
        """深度对比两次未读数统计是否有变化"""
        if old_counts is None:
            return True  # 首次推送
        
        # 对比账户级别和文件夹级别未读数
        return old_counts != new_counts

    async def send_personal_message(self, message: str, user_id: int):
        if user_id in self.active_connections:
            await self.active_connections[user_id].send_text(message)

    async def broadcast_unread_count_update(self, user_id: int):
        """手动触发未读数更新,对比后有变化才推送"""
        if user_id in self.active_connections:
            db = SessionLocal()
            try:
                counts = get_unread_counts(db, user_id)
                last_counts = self.last_unread_counts.get(user_id)
                
                if self._counts_changed(last_counts, counts):
                    await self.send_personal_message(
                        json.dumps({"status": "unread_count_update", "counts": counts}),
                        user_id
                    )
                    self.last_unread_counts[user_id] = counts
            finally:
                db.close()

    async def broadcast_sync_failed_update(self, user_id: int):
        """手动触发同步失败状态更新，每次都查询数据库并推送"""
        if user_id in self.active_connections:
            db = SessionLocal()
            try:
                failed_accounts = get_sync_failed_accounts(db, user_id)
                # 每次都推送，不使用缓存
                await self.send_personal_message(
                    json.dumps({"status": "sync_failed_update", "failed_accounts": failed_accounts}),
                    user_id
                )
            finally:
                db.close()

manager = ConnectionManager()

def get_sync_failed_accounts(db: Session, user_id: int) -> list:
    """
    获取用户同步失败的邮箱账户列表
    每次都查询数据库，不使用缓存
    """
    failed_accounts = db.query(EmailAccount).filter(
        EmailAccount.user_id == user_id,
        EmailAccount.sync_status == "failed"
    ).all()
    
    return [{"id": account.id, "name": account.name or account.email} for account in failed_accounts]

def get_unread_counts(db: Session, user_id: int) -> dict:
    """
    获取用户所有邮箱账户的未读邮件数，包括账户级别和文件夹级别
    【v2.0改造】使用EmailFolder关联表统计
    返回每个文件夹自己的未读数（不包括子文件夹），累加逻辑由前端处理
    """
    accounts = db.query(EmailAccount).filter(EmailAccount.user_id == user_id).all()
    result = {
        "accounts": {},  # 账户级别未读数
        "folders": {}    # 文件夹级别未读数（原始数据，不累加）
    }
    
    for account in accounts:
        # 账户级别未读数（去重统计）
        account_count = db.query(func.count(func.distinct(Email.id))).join(
            EmailFolder, Email.id == EmailFolder.email_id
        ).filter(
            Email.email_account_id == account.id,
            EmailFolder.email_account_id == account.id,
            Email.is_read == False
        ).scalar()
        
        result["accounts"][account.id] = account_count or 0
        
        # 文件夹级别未读数（使用EmailFolder关联表）
        folder_counts = db.query(
            EmailFolder.folder_name,
            func.count(func.distinct(EmailFolder.email_id)).label('count')
        ).join(
            Email, Email.id == EmailFolder.email_id
        ).filter(
            EmailFolder.email_account_id == account.id,
            Email.email_account_id == account.id,
            Email.is_read == False
        ).group_by(
            EmailFolder.folder_name
        ).all()
        
        if account.id not in result["folders"]:
            result["folders"][account.id] = {}
        
        # 存储每个文件夹自己的未读数（不累加子文件夹）
        for folder_name, count in folder_counts:
            result["folders"][account.id][folder_name] = count
        
        # 添加UNREAD虚拟文件夹，统计所有文件夹的未读邮件（包括Junk等）
        # 这样可以确保垃圾邮件等所有文件夹的未读数都会被统计到UNREAD中
        unread_total = db.query(func.count(func.distinct(Email.id))).join(
            EmailFolder, Email.id == EmailFolder.email_id
        ).filter(
            Email.email_account_id == account.id,
            EmailFolder.email_account_id == account.id,
            Email.is_read == False
        ).scalar()
        
        result["folders"][account.id]["UNREAD"] = unread_total or 0
        
    return result

async def unread_count_updater(user_id: int):
    """
    每5秒检查未读邮件数变化,仅在变化时推送
    """
    while True:
        await asyncio.sleep(5)
        db = SessionLocal()
        try:
            counts = get_unread_counts(db, user_id)
            last_counts = manager.last_unread_counts.get(user_id)
            
            # 对比是否有变化,有变化才推送
            if manager._counts_changed(last_counts, counts):
                await manager.send_personal_message(
                    json.dumps({"status": "unread_count_update", "counts": counts}),
                    user_id
                )
                manager.last_unread_counts[user_id] = counts
        except Exception as e:
            logger.error(f"更新未读邮件数时出错: {e}")
        finally:
            db.close()

async def sync_failed_updater(user_id: int):
    """
    每1分钟检查同步失败状态,每次都推送
    """
    while True:
        await asyncio.sleep(60)
        try:
            await manager.broadcast_sync_failed_update(user_id)
        except Exception as e:
            logger.error(f"更新同步失败状态时出错: {e}")


async def token_keeper(user_id: int, initial_token: str):
    """
    每30秒检查token是否达到50%时长，达到则推送retoken消息
    """
    token = initial_token
    
    while True:
        await asyncio.sleep(30)  # 30秒检查一次
        
        try:
            # 解码token获取过期时间
            payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
            exp = payload.get('exp')
            username = payload.get('sub')
            
            if not exp or not username:
                continue
            
            # 计算剩余有效时间
            now = datetime.utcnow()
            exp_datetime = datetime.utcfromtimestamp(exp)
            remaining_time = (exp_datetime - now).total_seconds()
            
            # 计算总有效时间和50%阈值
            total_time = ACCESS_TOKEN_EXPIRE_MINUTES * 60
            threshold_time = total_time * 0.5
            
            remaining_hours = remaining_time / 3600
            remaining_minutes = remaining_time / 60
            
            # 如果剩余时间少于50%且token仍然有效，生成新token并推送
            if 0 < remaining_time < threshold_time:
                # 生成新token
                new_expire = datetime.utcnow() + timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES)
                new_payload = {
                    "sub": username,
                    "exp": new_expire
                }
                new_token = jwt.encode(new_payload, SECRET_KEY, algorithm=ALGORITHM)
                
                # 推送retoken消息
                await manager.send_personal_message(
                    json.dumps({
                        "status": "retoken",
                        "token": new_token
                    }),
                    user_id
                )
                
                # 记录续期日志
                logger.info(f"Token保活推送: 用户={username}, 用户ID={user_id}, 剩余时间={remaining_minutes:.2f}分钟")
                
                # 更新token变量以便下次检查
                token = new_token
                logger.info(f"Token保活任务内部token已更新: 用户ID={user_id}")
            elif remaining_time <= 0:
                break
                
        except jwt.ExpiredSignatureError:
            break
        except jwt.JWTError as e:
            logger.error(f"用户 {user_id} 的token解码失败: {str(e)}")
            break
        except Exception as e:
            logger.error(f"Token保活检查错误 (用户 {user_id}): {str(e)}", exc_info=True)


@router.websocket("/ws/sync")
async def websocket_endpoint(
    websocket: WebSocket,
    token: str = Query(...),
):
    """
    WebSocket端点，用于实时邮件同步通知
    """
    # 在这里使用一次性的db session来验证token
    db = SessionLocal()
    try:
        user = get_current_user_from_token(token, db)
        if not user:
            await websocket.close(code=1008)
            return
        user_id = user.id
    finally:
        db.close()

    await manager.connect(websocket, user_id)

    # 立即发送一次未读邮件数统计(首次连接直接推送)
    db = SessionLocal()
    try:
        counts = get_unread_counts(db, user_id)
        await manager.send_personal_message(
            json.dumps({"status": "unread_count_update", "counts": counts}),
            user_id
        )
        manager.last_unread_counts[user_id] = counts  # 保存首次推送数据
    except Exception as e:
        logger.error(f"发送初始未读邮件数时出错: {e}")
    finally:
        db.close()

    # 立即发送一次同步失败状态(首次连接直接推送)
    db = SessionLocal()
    try:
        failed_accounts = get_sync_failed_accounts(db, user_id)
        await manager.send_personal_message(
            json.dumps({"status": "sync_failed_update", "failed_accounts": failed_accounts}),
            user_id
        )
    except Exception as e:
        logger.error(f"发送初始同步失败状态时出错: {e}")
    finally:
        db.close()

    # 启动未读邮件数更新任务（每5秒更新一次）
    unread_updater_task = asyncio.create_task(unread_count_updater(user_id))
    
    # 启动同步失败状态更新任务（每5秒更新一次）
    sync_failed_updater_task = asyncio.create_task(sync_failed_updater(user_id))
    
    # 启动token保活任务（每5分钟检查一次）
    token_keeper_task = asyncio.create_task(token_keeper(user_id, token))

    try:
        while True:
            data = await websocket.receive_text()
            message = json.loads(data)

            # 处理客户端发送的消息
            if isinstance(message, dict):
                action = message.get("action")
                
                if action == "refresh_unread_counts":
                    # 客户端请求刷新未读数，直接获取并推送，绕过缓存比较
                    db_session = SessionLocal()
                    try:
                        counts = get_unread_counts(db_session, user_id)
                        await manager.send_personal_message(
                            json.dumps({"status": "unread_count_update", "counts": counts}),
                            user_id
                        )
                        # 更新缓存，以便下一次自动轮询时比较
                        manager.last_unread_counts[user_id] = counts
                    except Exception as e:
                        logger.error(f"处理 refresh_unread_counts 时出错: {e}", exc_info=True)
                    finally:
                        db_session.close()
                elif action == "refresh_sync_failed":
                    # 客户端请求刷新同步失败状态，直接获取并推送
                    try:
                        await manager.broadcast_sync_failed_update(user_id)
                    except Exception as e:
                        logger.error(f"处理 refresh_sync_failed 时出错: {e}", exc_info=True)

    except WebSocketDisconnect:
        manager.disconnect(user_id)
        unread_updater_task.cancel()
        sync_failed_updater_task.cancel()
        token_keeper_task.cancel()
